{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b40296a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import numpy as np \n",
    "import matplotlib.pyplot as plt\n",
    "import warnings\n",
    "import matplotlib\n",
    "import random\n",
    "import seaborn as sns\n",
    "import time\n",
    "from sklift.metrics import (qini_auc_score)\n",
    "plt.style.use('fivethirtyeight')\n",
    "matplotlib.rcParams['figure.figsize'] = [8, 8]\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# logging.basicConfig(level=logging.INFO)\n",
    "pd.options.display.float_format = '{:.4f}'.format\n",
    "sns.set(rc={'figure.figsize':(5.7,4.27)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9e67e03c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 评估函数\n",
    "def evaluate_bin(t, yf, tau_true, tau_pred):\n",
    "    pehe = np.sqrt(np.mean(np.square(tau_pred - tau_true)))  # PEHE error\n",
    "\n",
    "    ate_pred = np.mean(tau_pred)\n",
    "    atc_pred = np.mean(tau_pred[(1 - t) > 0])\n",
    "    att_pred = np.mean(tau_pred[t > 0])\n",
    "\n",
    "    att = np.mean(tau_true[t > 0])\n",
    "    ate = np.mean(tau_true)\n",
    "\n",
    "    bias_att = np.abs(att_pred - att)  # the error of att\n",
    "    bias_ate = np.abs(ate_pred - ate)  # the error of ate\n",
    "\n",
    "    return {\"E_pehe\": pehe, \"E_att\": bias_att, \"E_ate\": bias_ate}\n",
    "\n",
    "\n",
    "def eval_mu(y_real, t, yhat_0, yhat_1):\n",
    "    auc_c = roc_auc_score(y_real, yhat_0, sample_weight = 1-t)\n",
    "    auc_t = roc_auc_score(y_real, yhat_1, sample_weight = t)\n",
    "    return {\"auc_c\":auc_c, \"auc_t\":auc_t}\n",
    "\n",
    "def eval_auuc(y_real, pred_tau, t_real):\n",
    "    return qini_auc_score(y_real, pred_tau, t_real)\n",
    "\n",
    "\n",
    "# 加载并解析\n",
    "def load_data(file_path):\n",
    "    \"\"\" Load data set \"\"\"\n",
    "    data_in = np.load(file_path)\n",
    "    data = {'x': data_in['x'], 't': data_in['t'], 'yf': data_in['yf']}\n",
    "    try:\n",
    "        data['ycf'] = data_in['ycf']\n",
    "        data[\"mu0\"] = data_in['mu0']\n",
    "        data[\"mu1\"] = data_in['mu1']\n",
    "    except:\n",
    "        data['ycf'] = None\n",
    "\n",
    "    try:\n",
    "        data['e'] = data_in['e']\n",
    "    except:\n",
    "        data['e'] = np.ones_like(data_in['yf'])\n",
    "\n",
    "    try:\n",
    "        data['tau'] = data_in['tau']\n",
    "        data['IS_SYNT'] = True\n",
    "    except:\n",
    "        data['tau'] = np.array([None])\n",
    "        data['IS_SYNT'] = False\n",
    "\n",
    "\n",
    "    data['dim'] = data['x'].shape[1] # 特征维度\n",
    "    data['n'] = data['x'].shape[0] # 样本数\n",
    "\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "83bd9b3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(834003, 83, 5)\n",
      "(181669, 83, 5)\n",
      "(834003, 5)\n",
      "(181669, 5)\n"
     ]
    }
   ],
   "source": [
    "name = \"lzd\"\n",
    "dd = \"real\"\n",
    "\n",
    "train_data_path = \"./uplift_data/lzd_real_data_id_v0/real_bin_set_full.5.train.npz\".format(str(name+dd))\n",
    "test_data_path = \"./uplift_data/lzd_real_data_id_v0/real_bin_set_full.5.test.npz\".format(str(name+dd))\n",
    "\n",
    "dict_train = load_data(train_data_path)\n",
    "dict_test = load_data(test_data_path)\n",
    "\n",
    "\n",
    "print(np.array(dict_train[\"x\"]).shape)\n",
    "print(np.array(dict_test[\"x\"]).shape)\n",
    "\n",
    "print(np.array(dict_train[\"yf\"]).shape)\n",
    "print(np.array(dict_test[\"yf\"]).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9eab5b6d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running... 0\n",
      "running... 1\n",
      "running... 2\n",
      "running... 3\n",
      "running... 4\n",
      "testing eval results:\n",
      "E_pehe: nan +/- nan\n",
      "E_att: 0.012280 +/- 0.000250\n",
      "E_ate: nan +/- nan\n",
      "AUUC: 0.013207 +/- 0.000848\n"
     ]
    }
   ],
   "source": [
    "from causalml.inference.tree import UpliftTreeClassifier, UpliftRandomForestClassifier\n",
    "from econml.dml import CausalForestDML\n",
    "    \n",
    "# save for predictions\n",
    "train_eval_result={\"E_pehe\":[], \"E_att\":[], \"E_ate\":[], \"AUUC\":[]}\n",
    "test_eval_result={\"E_pehe\":[], \"E_att\":[], \"E_ate\":[], \"AUUC\":[]}\n",
    "test_random_eval_result={\"E_pehe\":[], \"E_att\":[], \"E_ate\":[], \"AUUC\":[]}\n",
    "\n",
    "num_exp = dict_train[\"x\"].shape[2]\n",
    "\n",
    "include_random=False\n",
    "\n",
    "\n",
    "\n",
    "# 模拟数据重复生成了多份数据，每一份数据对应一次重复实验\n",
    "for i_exp in range(num_exp):\n",
    "    print(\"running... \" + str(i_exp))\n",
    "    start = time.time()\n",
    "    # train\n",
    "    X_train = dict_train[\"x\"][:,:,i_exp]\n",
    "    yf_train = dict_train[\"yf\"][:,i_exp]\n",
    "    t_train = dict_train[\"t\"][:,i_exp]\n",
    "\n",
    "    uplift_model = UpliftRandomForestClassifier(control_name=\"0\",\n",
    "                                                max_features=10,\n",
    "                                                min_samples_leaf=100,\n",
    "                                                random_state=42,\n",
    "                                                n_estimators=50,\n",
    "                                                max_depth=5,\n",
    "                                                n_jobs=8)\n",
    "    uplift_model.fit(X_train,\n",
    "                     treatment=t_train.astype(str),\n",
    "                     y=yf_train)\n",
    "    # Use defaults\n",
    "    est = CausalForestDML()\n",
    "    # Or specify hyperparameters\n",
    "    # 200, 10, 0.1, 0.5，5\n",
    "    est = CausalForestDML(n_estimators=200,       \n",
    "                          max_depth=10,\n",
    "                          max_features=0.1,\n",
    "                          max_samples=0.5,\n",
    "                          subforest_size=10,\n",
    "                          random_state=42)\n",
    "    est.fit(yf_train, t_train, X=X_train)\n",
    "\n",
    "    # predict\n",
    "    X_test = dict_test[\"x\"][:,:, i_exp]\n",
    "    yf_test = dict_test[\"yf\"][:, i_exp]\n",
    "    t_test = dict_test[\"t\"][:, i_exp]\n",
    "\n",
    "    pred_result = est.effect(X_test)\n",
    "    _auuc = eval_auuc(yf_test.reshape(-1), pred_result.reshape(-1), t_test.reshape(-1))\n",
    "\n",
    "    att = np.mean(yf_test.reshape(-1)[t_test.reshape(-1)==1]) - np.mean(yf_test.reshape(-1)[t_test.reshape(-1)==0]) \n",
    "    pred_att = np.mean(pred_result.reshape(-1)[t_test.reshape(-1)==1])\n",
    "    E_att = np.abs(pred_att-att)\n",
    "\n",
    "    test_eval_result[\"E_att\"].append(E_att)\n",
    "    test_eval_result[\"AUUC\"].append(_auuc)\n",
    "\n",
    "\n",
    "print(\"testing eval results:\")\n",
    "for k in test_eval_result.keys():\n",
    "    val = np.mean(test_eval_result[k])\n",
    "    num_exp = len(test_eval_result[k])\n",
    "    std = np.std(test_eval_result[k])/np.sqrt(num_exp)\n",
    "    print(k+\": %.6f\"%val + \" +/- %.6f\"%std )\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87614a34",
   "metadata": {},
   "source": [
    "# BART\n",
    "https://github.com/JakeColtman/bartpy\n",
    "pip install git+https://github.com/JakeColtman/bartpy.git --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "591f2390",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bartpy.sklearnmodel import SklearnModel\n",
    "\n",
    "test_eval_result={\"E_pehe\":[], \"E_att\":[], \"E_ate\":[], \"AUUC\":[]}\n",
    "\n",
    "for i_exp in range(num_exp):\n",
    "    print(\"running... \" + str(i_exp))\n",
    "    start = time.time()\n",
    "    # train\n",
    "    X_train = dict_train[\"x\"][:, :, i_exp]\n",
    "    yf_train = dict_train[\"yf\"][:, i_exp]\n",
    "    t_train = dict_train[\"t\"][:, i_exp]\n",
    "\n",
    "    bart_model = SklearnModel(n_trees=20, n_jobs=1)  # Use default parameters\n",
    "    bart_model.fit(X_train, yf_train)  # Fit the model\n",
    " \n",
    "\n",
    "    # predict\n",
    "    X_test = dict_test[\"x\"][:, :, i_exp]\n",
    "    yf_test = dict_test[\"yf\"][:, i_exp]\n",
    "    t_test = dict_test[\"t\"][:, i_exp]\n",
    "\n",
    "    pred_result = bart_model.predict(X_test)\n",
    "   \n",
    "\n",
    "    _auuc = qini_auc_score(yf_test.reshape(-1), pred_result, t_test.reshape(-1))\n",
    "    att = np.mean(yf_test.reshape(-1)[t_test.reshape(-1)==1]) - np.mean(yf_test.reshape(-1)[t_test.reshape(-1)==0]) \n",
    "    pred_att = np.mean(pred_result[t_test.reshape(-1)==1])\n",
    "    E_att = np.abs(pred_att-att)\n",
    "    \n",
    "    test_eval_result[\"E_att\"].append(E_att)\n",
    "    test_eval_result[\"AUUC\"].append(_auuc)\n",
    "\n",
    "\n",
    "print(\"testing eval results:\")\n",
    "for k in test_eval_result.keys():\n",
    "    val = np.mean(test_eval_result[k])\n",
    "    num_exp = len(test_eval_result[k])\n",
    "    std = np.std(test_eval_result[k])/np.sqrt(num_exp)\n",
    "    print(k+\": %.6f\"%val + \" +/- %.6f\"%std )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faba73d2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
